import gzip
import json
import pickle
from typing import Any, List, Union, Optional, Tuple

import numpy as np
try:
    from dtw import dtw  # type: ignore
except Exception:
    dtw = None  # type: ignore
try:
    from fastdtw import fastdtw  # type: ignore
except Exception:
    fastdtw = None  # type: ignore
from habitat.core.dataset import Episode
from habitat.core.embodied_task import Action, EmbodiedTask, Measure
from habitat.core.logging import logger
from habitat.core.registry import registry
from habitat.core.simulator import Simulator
from habitat.core.utils import try_cv2_import
from habitat.tasks.nav.nav import DistanceToGoal, Success
from habitat.tasks.utils import cartesian_to_polar
from habitat.utils.geometry_utils import quaternion_rotate_vector
from habitat.utils.visualizations import fog_of_war
from habitat.utils.visualizations import maps as habitat_maps
from numpy import ndarray
from omegaconf import DictConfig
from . import maps
from .rxr_dataset import RxRVLNDatasetV1

cv2 = try_cv2_import()


def euclidean_distance(
    pos_a: Union[List[float], ndarray], pos_b: Union[List[float], ndarray]
) -> float:
    return np.linalg.norm(np.array(pos_b) - np.array(pos_a), ord=2)


def _dtw_fallback(seq_a, seq_b, dist):
    """Simple DTW fallback returning (distance, None, None, None).

    Implemented to avoid hard dependency on dtw/fastdtw. Complexity O(n*m).
    """
    n, m = len(seq_a), len(seq_b)
    cost = np.full((n + 1, m + 1), np.inf, dtype=np.float32)
    cost[0, 0] = 0.0
    for i in range(1, n + 1):
        a_i = seq_a[i - 1]
        for j in range(1, m + 1):
            b_j = seq_b[j - 1]
            d = dist(a_i, b_j)
            cost[i, j] = d + min(cost[i - 1, j], cost[i, j - 1], cost[i - 1, j - 1])
    return float(cost[n, m]), None, None, None


@registry.register_measure
class PathLength(Measure):
    """Path Length (PL)
    PL = sum(geodesic_distance(agent_prev_position, agent_position)
            over all agent positions.
    """

    cls_uuid: str = "path_length"

    def __init__(self, sim: Simulator, *args: Any, **kwargs: Any):
        self._sim = sim
        super().__init__(**kwargs)

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def reset_metric(self, *args: Any, **kwargs: Any):
        self._previous_position = self._sim.get_agent_state().position
        self._metric = 0.0

    def update_metric(self, *args: Any, **kwargs: Any):
        current_position = self._sim.get_agent_state().position
        self._metric += euclidean_distance(
            current_position, self._previous_position
        )
        self._previous_position = current_position


@registry.register_measure
class OracleNavigationError(Measure):
    """Oracle Navigation Error (ONE)
    ONE = min(geosdesic_distance(agent_pos, goal)) over all points in the
    agent path.
    """

    cls_uuid: str = "oracle_navigation_error"

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def reset_metric(self, *args: Any, task: EmbodiedTask, **kwargs: Any):
        task.measurements.check_measure_dependencies(
            self.uuid, [DistanceToGoal.cls_uuid]
        )
        self._metric = float("inf")
        self.update_metric(task=task)

    def update_metric(self, *args: Any, task: EmbodiedTask, **kwargs: Any):
        distance_to_target = task.measurements.measures[
            DistanceToGoal.cls_uuid
        ].get_metric()
        self._metric = min(self._metric, distance_to_target)


@registry.register_measure
class OracleSuccess(Measure):
    """Oracle Success Rate (OSR). OSR = I(ONE <= goal_radius)"""

    cls_uuid: str = "oracle_success"

    def __init__(self, *args: Any, config: Any, **kwargs: Any):
        print(f"in oracle success init: args = {args}, kwargs = {kwargs}")
        self._config = config
        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def reset_metric(self, *args: Any, task: EmbodiedTask, **kwargs: Any):
        task.measurements.check_measure_dependencies(
            self.uuid, [DistanceToGoal.cls_uuid]
        )
        self._metric = 0.0
        self.update_metric(task=task)

    def update_metric(self, *args: Any, task: EmbodiedTask, **kwargs: Any):
        d = task.measurements.measures[DistanceToGoal.cls_uuid].get_metric()
        success_distance = getattr(self._config, "SUCCESS_DISTANCE", getattr(self._config, "success_distance", 3.0))
        self._metric = float(self._metric or d < success_distance)


@registry.register_measure
class OracleSPL(Measure):
    """OracleSPL (Oracle Success weighted by Path Length)
    OracleSPL = max(SPL) over all points in the agent path.
    """

    cls_uuid: str = "oracle_spl"

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def reset_metric(self, *args: Any, task: EmbodiedTask, **kwargs: Any):
        task.measurements.check_measure_dependencies(self.uuid, ["spl"])
        self._metric = 0.0

    def update_metric(self, *args: Any, task: EmbodiedTask, **kwargs: Any):
        spl = task.measurements.measures["spl"].get_metric()
        self._metric = max(self._metric, spl)

@registry.register_measure
class PL(Measure):
    r"""SPL (Success weighted by Path Length)

    ref: On Evaluation of Embodied Agents - Anderson et. al
    https://arxiv.org/pdf/1807.06757.pdf
    The measure depends on Distance to Goal measure and Success measure
    to improve computational
    performance for sophisticated goal areas.
    """

    def __init__(
        self, sim: Simulator, config: "DictConfig", *args: Any, **kwargs: Any
    ):
        self._previous_position: Union[None, np.ndarray, List[float]] = None
        self._start_end_episode_distance: Optional[float] = None
        self._agent_episode_distance: Optional[float] = None
        self._episode_view_points: Optional[
            List[Tuple[float, float, float]]
        ] = None
        self._sim = sim
        self._config = config

        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return "pl"

    def reset_metric(self, episode, task, *args: Any, **kwargs: Any):
        task.measurements.check_measure_dependencies(
            self.uuid, [DistanceToGoal.cls_uuid, Success.cls_uuid]
        )

        self._previous_position = self._sim.get_agent_state().position
        self._agent_episode_distance = 0.0
        self._start_end_episode_distance = task.measurements.measures[
            DistanceToGoal.cls_uuid
        ].get_metric()
        self.update_metric(  # type:ignore
            episode=episode, task=task, *args, **kwargs
        )

    def _euclidean_distance(self, position_a, position_b):
        return np.linalg.norm(position_b - position_a, ord=2)

    def update_metric(
        self, episode, task: EmbodiedTask, *args: Any, **kwargs: Any
    ):
        # ep_success = task.measurements.measures[Success.cls_uuid].get_metric()

        current_position = self._sim.get_agent_state().position
        self._agent_episode_distance += self._euclidean_distance(
            current_position, self._previous_position
        )

        self._previous_position = current_position

        self._metric = (
            self._start_end_episode_distance
            / max(
                self._start_end_episode_distance, self._agent_episode_distance
            )
        )

@registry.register_measure
class StepsTaken(Measure):
    """Counts the number of times update_metric() is called. This is equal to
    the number of times that the agent takes an action. STOP counts as an
    action.
    """

    cls_uuid: str = "steps_taken"

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def reset_metric(self, *args: Any, **kwargs: Any):
        self._metric = 0.0

    def update_metric(self, *args: Any, **kwargs: Any):
        self._metric += 1.0


@registry.register_measure
class WaypointRewardMeasure(Measure):
    """A reward measure used for training VLN-CE agents via RL."""

    def __init__(self, *args: Any, sim: Simulator, config: Any, **kwargs: Any) -> None:
        self._sim = sim
        self._slack_reward = config.slack_reward
        self._use_distance_scaled_slack_reward = config.use_distance_scaled_slack_reward
        self._scale_slack_on_prediction = config.scale_slack_on_prediction
        self._success_reward = config.success_reward
        self._distance_scalar = config.distance_scalar
        self._prev_position = None
        super().__init__()

    def reset_metric(self, *args: Any, task: EmbodiedTask, **kwargs: Any) -> None:
        task.measurements.check_measure_dependencies(self.uuid, [DistanceToGoal.cls_uuid, Success.cls_uuid])
        self._previous_distance_to_goal = task.measurements.measures["distance_to_goal"].get_metric()
        self._metric = 0.0
        self._prev_position = np.take(self._sim.get_agent_state().position, [0, 2])

    def _get_scaled_slack_reward(self, action: Action) -> float:
        if isinstance(action["action"], int):
            return self._slack_reward

        if not self._use_distance_scaled_slack_reward:
            return self._slack_reward

        agent_pos = np.take(self._sim.get_agent_state().position, [0, 2])
        slack_distance = (
            action["action_args"]["r"]
            if self._scale_slack_on_prediction and action["action"] != "STOP"
            else np.linalg.norm(self._prev_position - agent_pos)
        )
        scaled_slack_reward = self._slack_reward * slack_distance / 0.25
        self._prev_position = agent_pos
        return min(self._slack_reward, scaled_slack_reward)

    def _progress_to_goal(self, task: EmbodiedTask) -> float:
        distance_to_goal = task.measurements.measures["distance_to_goal"].get_metric()
        distance_to_goal_delta = self._previous_distance_to_goal - distance_to_goal
        if np.isnan(distance_to_goal_delta) or np.isinf(distance_to_goal_delta):
            l = self._sim.get_agent_state().position
            logger.error(
                f"\nNaN or inf encountered in distance measure. agent location: {l}",
            )
            distance_to_goal_delta = -1.0
        self._previous_distance_to_goal = distance_to_goal
        return self._distance_scalar * distance_to_goal_delta

    def update_metric(self, *args: Any, action: Action, task: EmbodiedTask, **kwargs: Any) -> None:
        reward = self._get_scaled_slack_reward(action)
        reward += self._progress_to_goal(task)
        reward += self._success_reward * task.measurements.measures["success"].get_metric()
        self._metric = reward

    @staticmethod
    def _get_uuid(*args: Any, **kwargs: Any) -> str:
        return "waypoint_reward_measure"


@registry.register_measure
class NDTW(Measure):
    """NDTW (Normalized Dynamic Time Warping)
    ref: https://arxiv.org/abs/1907.05446
    """

    cls_uuid: str = "ndtw"

    def __init__(self, *args: Any, sim: Simulator, config: Any, **kwargs: Any):
        self._sim = sim
        self._config = config
        # Choose DTW implementation with graceful fallback
        if getattr(config, "FDTW", False) and (globals().get("fastdtw") is not None):
            self.dtw_func = fastdtw
        elif globals().get("dtw") is not None:
            self.dtw_func = dtw
        else:
            self.dtw_func = _dtw_fallback

        if "{role}" in config.GT_PATH:
            self.gt_json = {}
            for role in RxRVLNCEDatasetV1.annotation_roles:
                with gzip.open(config.GT_PATH.format(split=config.SPLIT, role=role), "rt") as f:
                    self.gt_json.update(json.load(f))
        else:
            with gzip.open(config.GT_PATH.format(split=config.SPLIT), "rt") as f:
                self.gt_json = json.load(f)

        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def reset_metric(self, *args: Any, episode, **kwargs: Any):
        self.locations = []
        self.gt_locations = self.gt_json[episode.episode_id]["locations"]
        self.update_metric()

    def update_metric(self, *args: Any, **kwargs: Any):
        current_position = self._sim.get_agent_state().position.tolist()
        if len(self.locations) == 0:
            self.locations.append(current_position)
        else:
            if current_position == self.locations[-1]:
                return
            self.locations.append(current_position)

        dtw_result = self.dtw_func(self.locations, self.gt_locations, dist=euclidean_distance)
        dtw_distance = dtw_result[0] if isinstance(dtw_result, (list, tuple)) else float(dtw_result)

        nDTW = np.exp(-dtw_distance / (len(self.gt_locations) * self._config.SUCCESS_DISTANCE))
        self._metric = float(nDTW)


@registry.register_measure
class SDTW(Measure):
    """SDTW (Success Weighted be nDTW)
    ref: https://arxiv.org/abs/1907.05446
    """

    cls_uuid: str = "sdtw"

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def reset_metric(self, *args: Any, task: EmbodiedTask, **kwargs: Any):
        task.measurements.check_measure_dependencies(self.uuid, [NDTW.cls_uuid, Success.cls_uuid])
        self.update_metric(task=task)

    def update_metric(self, *args: Any, task: EmbodiedTask, **kwargs: Any):
        ep_success = task.measurements.measures[Success.cls_uuid].get_metric()
        nDTW = task.measurements.measures[NDTW.cls_uuid].get_metric()
        self._metric = float(ep_success * nDTW)


@registry.register_measure
class TopDownMapVLNCE(Measure):
    """A top down map that optionally shows VLN-related visual information
    such as MP3D node locations and MP3D agent traversals.
    """

    cls_uuid: str = "top_down_map_vlnce"

    def __init__(self, *args: Any, sim: Simulator, config: Any, **kwargs: Any) -> None:
        self._sim = sim
        self._config = config
        self._step_count = None
        self._map_resolution = getattr(config, "MAP_RESOLUTION", 1024)
        self._previous_xy_location = None
        self._top_down_map = None
        self._meters_per_pixel = None
        self.current_node = ""
        with open(self._config.GRAPHS_FILE, "rb") as f:
            self._conn_graphs = pickle.load(f)
        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def get_original_map(self) -> ndarray:
        top_down_map = habitat_maps.get_topdown_map_from_sim(
            self._sim,
            map_resolution=self._map_resolution,
            draw_border=getattr(self._config, "DRAW_BORDER", True),
            meters_per_pixel=self._meters_per_pixel,
        )

        self._fog_of_war_mask = None
        if getattr(self._config.FOG_OF_WAR, "DRAW", False):
            self._fog_of_war_mask = np.zeros_like(top_down_map)

        return top_down_map

    def reset_metric(self, *args: Any, episode: Episode, **kwargs: Any) -> None:
        self._scene_id = episode.scene_id.split("/")[-2]
        self._step_count = 0
        self._metric = None
        self._meters_per_pixel = habitat_maps.calculate_meters_per_pixel(
            self._map_resolution, self._sim
        )
        self._top_down_map = self.get_original_map()
        agent_position = self._sim.get_agent_state().position
        scene_id = episode.scene_id.split("/")[-1].split(".")[0]
        a_x, a_y = habitat_maps.to_grid(
            agent_position[2],
            agent_position[0],
            self._top_down_map.shape[0:2],
            sim=self._sim,
        )
        self._previous_xy_location = (a_y, a_x)

        if getattr(self._config.FOG_OF_WAR, "DRAW", False):
            self._fog_of_war_mask = fog_of_war.reveal_fog_of_war(
                self._top_down_map,
                self._fog_of_war_mask,
                np.array([a_x, a_y]),
                self.get_polar_angle(),
                fov=getattr(self._config.FOG_OF_WAR, "FOV", 90),
                max_line_len=getattr(self._config.FOG_OF_WAR, "VISIBILITY_DIST", 5.0)
                / habitat_maps.calculate_meters_per_pixel(
                    self._map_resolution, sim=self._sim
                ),
            )

        if getattr(self._config, "DRAW_FIXED_WAYPOINTS", False):
            maps.draw_mp3d_nodes(
                self._top_down_map,
                self._sim,
                episode,
                self._conn_graphs[scene_id],
                self._meters_per_pixel,
            )

        if getattr(self._config, "DRAW_SHORTEST_PATH", False):
            shortest_path_points = self._sim.get_straight_shortest_path_points(
                agent_position, episode.goals[0].position
            )
            maps.draw_straight_shortest_path_points(
                self._top_down_map,
                self._sim,
                self._map_resolution,
                shortest_path_points,
            )

        if getattr(self._config, "DRAW_REFERENCE_PATH", False):
            maps.draw_reference_path(
                self._top_down_map,
                self._sim,
                episode,
                self._map_resolution,
                self._meters_per_pixel,
            )

        if getattr(self._config, "DRAW_SOURCE_AND_TARGET", False):
            maps.draw_source_and_target(
                self._top_down_map,
                self._sim,
                episode,
                self._meters_per_pixel,
            )

        self._nearest_node = maps.get_nearest_node(
            self._conn_graphs[scene_id], np.take(agent_position, (0, 2))
        )
        nn_position = self._conn_graphs[self._scene_id].nodes[
            self._nearest_node
        ]["position"]
        self.s_x, self.s_y = habitat_maps.to_grid(
            nn_position[2],
            nn_position[0],
            self._top_down_map.shape[0:2],
            self._sim,
        )
        self.update_metric()

    def update_metric(self, *args: Any, **kwargs: Any) -> None:
        self._step_count += 1
        (
            house_map,
            map_agent_pos,
        ) = self.update_map(self._sim.get_agent_state().position)

        self._metric = {
            "map": house_map,
            "fog_of_war_mask": self._fog_of_war_mask,
            "agent_map_coord": map_agent_pos,
            "agent_angle": self.get_polar_angle(),
            "bounds": {
                k: v
                for k, v in zip(
                    ["lower", "upper"],
                    self._sim.pathfinder.get_bounds(),
                )
            },
            "meters_per_px": self._meters_per_pixel,
        }

    def get_polar_angle(self) -> float:
        agent_state = self._sim.get_agent_state()
        ref_rotation = agent_state.rotation
        heading_vector = quaternion_rotate_vector(
            ref_rotation.inverse(), np.array([0, 0, -1])
        )
        phi = cartesian_to_polar(-heading_vector[2], heading_vector[0])[1]
        z_neg_z_flip = np.pi
        return np.array(phi) + z_neg_z_flip

    def update_map(self, agent_position: List[float]) -> None:
        a_x, a_y = habitat_maps.to_grid(
            agent_position[2],
            agent_position[0],
            self._top_down_map.shape[0:2],
            self._sim,
        )
        gradient_color = 15 + min(
            self._step_count * 245 // getattr(self._config, "MAX_EPISODE_STEPS", 500), 245
        )
        if self._top_down_map[a_x, a_y] != maps.MAP_SOURCE_POINT_INDICATOR:
            maps.drawline(
                self._top_down_map,
                self._previous_xy_location,
                (a_y, a_x),
                gradient_color,
                thickness=int(
                    self._map_resolution * 1.4 / maps.MAP_THICKNESS_SCALAR
                ),
                style="filled",
            )

        if getattr(self._config.FOG_OF_WAR, "DRAW", False):
            self._fog_of_war_mask = fog_of_war.reveal_fog_of_war(
                self._top_down_map,
                self._fog_of_war_mask,
                np.array([a_x, a_y]),
                self.get_polar_angle(),
                getattr(self._config.FOG_OF_WAR, "FOV", 90),
                max_line_len=getattr(self._config.FOG_OF_WAR, "VISIBILITY_DIST", 5.0)
                / habitat_maps.calculate_meters_per_pixel(
                    self._map_resolution, sim=self._sim
                ),
            )

        point_padding = int(0.2 / self._meters_per_pixel)
        prev_nearest_node = self._nearest_node
        self._nearest_node = maps.update_nearest_node(
            self._conn_graphs[self._scene_id],
            self._nearest_node,
            np.take(agent_position, (0, 2)),
        )
        if (
            self._nearest_node != prev_nearest_node
            and getattr(self._config, "DRAW_MP3D_AGENT_PATH", False)
        ):
            nn_position = self._conn_graphs[self._scene_id].nodes[
                self._nearest_node
            ]["position"]
            (prev_s_x, prev_s_y) = (self.s_x, self.s_y)
            self.s_x, self.s_y = habitat_maps.to_grid(
                nn_position[2],
                nn_position[0],
                self._top_down_map.shape[0:2],
                self._sim,
            )
            self._top_down_map[
                self.s_x - int(2.0 / 3.0 * point_padding) : self.s_x + int(2.0 / 3.0 * point_padding) + 1,
                self.s_y - int(2.0 / 3.0 * point_padding) : self.s_y + int(2.0 / 3.0 * point_padding) + 1,
            ] = gradient_color

            maps.drawline(
                self._top_down_map,
                (prev_s_y, prev_s_x),
                (self.s_y, self.s_x),
                gradient_color,
                thickness=int(
                    0.5 * np.round(
                        self._map_resolution / maps.MAP_THICKNESS_SCALAR
                    )
                ),
            )

        self._previous_xy_location = (a_y, a_x)
        map_agent_pos = (a_x, a_y)
        return self._top_down_map, map_agent_pos
